from flask import Flask, request, jsonify
import datetime
import tensorflow as tf
import numpy as np

app = Flask(__name__)

# 假设这是你的数据库
db = {
    "patients": {}
    # 其他集合...
}

# 已有的API接口...
@app.route("/api/patient/<case_id>/note", methods=["POST"])
def add_note(case_id):
    p = db["patients"].get(case_id)
    if not p:
        return jsonify({"ok": False, "message": "未找到该病例号"}), 404
    content = (request.json or {}).get("content", "").strip()
    if not content:
        return jsonify({"ok": False, "message": "内容不能为空"}), 400
    item = {
        "id": str(datetime.datetime.utcnow().timestamp()),
        "content": content,
        "createdAt": datetime.datetime.utcnow().isoformat() + "Z",
    }
    p["caseNotes"].insert(0, item)
    return jsonify({"ok": True, "data": item})

@app.route("/api/patient/<case_id>/report", methods=["POST"])
def add_report(case_id):
    p = db["patients"].get(case_id)
    if not p:
        return jsonify({"ok": False, "message": "未找到该病例号"}), 404
    body = request.json or {}
    title = body.get("title", "未命名报告")
    content = body.get("content", "").strip()
    item = {
        "id": str(datetime.datetime.utcnow().timestamp()),
        "title": title,
        "content": content,
        "createdAt": datetime.datetime.utcnow().isoformat() + "Z",
    }
    p["reports"].insert(0, item)
    return jsonify({"ok": True, "data": item})

@app.route("/api/patients", methods=["GET"])
def list_patients():
    return jsonify({"ok": True, "data": list(db["patients"].values())})

# 新添加的1: Fugl-Meyer (FMA) 量表自动评估接口
@app.route("/api/patient/<case_id>/fma-assessment", methods=["POST"])
def fma_assessment(case_id):
    """
    对患者进行Fugl-Meyer量表自动评估
    接收患者运动数据，返回评估结果
    """
    p = db["patients"].get(case_id)
    if not p:
        return jsonify({"ok": False, "message": "未找到该病例号"}), 404
    
    # 获取评估所需的运动数据
    assessment_data = request.json or {}
    movement_data = assessment_data.get("movementData")
    
    if not movement_data:
        return jsonify({"ok": False, "message": "运动数据不能为空"}), 400
    
    # 这里应该有实际的FMA评估逻辑
    # 简化示例：计算各项得分并汇总
    motor_function_score = calculate_motor_function(movement_data)
    balance_score = calculate_balance(movement_data)
    sensation_score = calculate_sensation(movement_data)
    
    total_score = motor_function_score + balance_score + sensation_score
    assessment_result = {
        "id": str(datetime.datetime.utcnow().timestamp()),
        "caseId": case_id,
        "motorFunctionScore": motor_function_score,
        "balanceScore": balance_score,
        "sensationScore": sensation_score,
        "totalScore": total_score,
        "assessmentDate": datetime.datetime.utcnow().isoformat() + "Z",
        # FMA总分最高为100分，根据总分判断恢复程度
        "recoveryLevel": get_recovery_level(total_score)
    }
    
    # 将评估结果保存到患者记录中
    if "fmaAssessments" not in p:
        p["fmaAssessments"] = []
    p["fmaAssessments"].insert(0, assessment_result)
    
    return jsonify({"ok": True, "data": assessment_result})

# 新添加的2: TensorFlow Lite模型推理接口
@app.route("/api/patient/<case_id>/tflite-inference", methods=["POST"])
def tflite_inference(case_id):
    """
    使用TensorFlow Lite模型对患者数据进行推理
    可用于辅助评估患者运动功能
    """
    p = db["patients"].get(case_id)
    if not p:
        return jsonify({"ok": False, "message": "未找到该病例号"}), 404
    
    # 获取推理所需的数据和模型信息
    inference_data = request.json or {}
    input_data = inference_data.get("inputData")
    model_name = inference_data.get("modelName", "default_model")
    
    if not input_data:
        return jsonify({"ok": False, "message": "输入数据不能为空"}), 400
    
    try:
        # 加载TensorFlow Lite模型并进行推理
        interpreter = tf.lite.Interpreter(model_path=f"models/{model_name}.tflite")
        interpreter.allocate_tensors()
        
        # 获取输入和输出张量
        input_details = interpreter.get_input_details()
        output_details = interpreter.get_output_details()
        
        # 准备输入数据
        input_array = np.array(input_data, dtype=np.float32).reshape(input_details[0]['shape'])
        interpreter.set_tensor(input_details[0]['index'], input_array)
        
        # 执行推理
        interpreter.invoke()
        
        # 获取推理结果
        output_data = interpreter.get_tensor(output_details[0]['index'])
        
        # 整理推理结果
        inference_result = {
            "id": str(datetime.datetime.utcnow().timestamp()),
            "caseId": case_id,
            "modelName": model_name,
            "result": output_data.tolist(),
            "inferenceDate": datetime.datetime.utcnow().isoformat() + "Z",
            "interpretation": interpret_inference_result(output_data)
        }
        
        # 保存推理结果
        if "aiAssessments" not in p:
            p["aiAssessments"] = []
        p["aiAssessments"].insert(0, inference_result)
        
        return jsonify({"ok": True, "data": inference_result})
        
    except Exception as e:
        return jsonify({"ok": False, "message": f"模型推理失败: {str(e)}"}), 500

# 辅助函数：计算运动功能得分
def calculate_motor_function(movement_data):
    # 实际应用中应该有更复杂的计算逻辑
    return min(66, max(0, len(movement_data) * 0.5))  # 示例计算

# 辅助函数：计算平衡功能得分
def calculate_balance(movement_data):
    # 实际应用中应该有更复杂的计算逻辑
    return min(14, max(0, len(movement_data) * 0.1))  # 示例计算

# 辅助函数：计算感觉功能得分
def calculate_sensation(movement_data):
    # 实际应用中应该有更复杂的计算逻辑
    return min(20, max(0, len(movement_data) * 0.2))  # 示例计算

# 辅助函数：根据总分判断恢复程度
def get_recovery_level(total_score):
    if total_score >= 90:
        return "几乎完全恢复"
    elif total_score >= 70:
        return "显著恢复"
    elif total_score >= 50:
        return "中度恢复"
    elif total_score >= 30:
        return "轻度恢复"
    else:
        return "恢复有限"

# 辅助函数：解释推理结果
def interpret_inference_result(output_data):
    # 根据模型输出提供可读性解释
    if np.max(output_data) > 0.7:
        return "运动功能评估显示良好，偏差在正常范围内"
    elif np.max(output_data) > 0.3:
        return "运动功能存在一定异常，建议进一步检查"
    else:
        return "运动功能明显异常，需要专业干预"

if __name__ == "__main__":
    app.run(host="127.0.0.1", port=5000, debug=True)
